#include "rbtree.h"

// 创建新的红黑树节点
rb_node * create_rb_node(int key) {
    rb_node *node = (rb_node *)malloc(sizeof(rb_node));
    if (node == NULL) {
        perror("内存分配失败");
        exit(EXIT_FAILURE);
    }
    node->key = key;
    node->color = RED;
    node->left = NULL;
    node->right = NULL;
    node->parent = NULL;
    return node;
}

// 创建红黑树
rb_tree * create_rb_tree() {
    rb_tree *tree = (rb_tree *)malloc(sizeof(rb_tree));
    if (tree == NULL) {
        perror("内存分配失败");
        exit(EXIT_FAILURE);
    }
    tree->nil = create_rb_node(0);
    tree->nil->color = BLACK;
    tree->root = tree->nil;
    return tree;
}

// 左旋操作
void left_rotate(rb_tree *tree, rb_node *node) {
    rb_node *right_child = node->right;
    node->right = right_child->left;
    if (right_child->left!= tree->nil) {
        right_child->left->parent = node;
    }
    right_child->parent = node->parent;
    if (node->parent == tree->nil) {
        tree->root = right_child;
    } else if (node == node->parent->left) {
        node->parent->left = right_child;
    } else {
        node->parent->right = right_child;
    }
    right_child->left = node;
    node->parent = right_child;
}

// 右旋操作
void right_rotate(rb_tree *tree, rb_node *node) {
    rb_node *left_child = node->left;
    node->left = left_child->right;
    if (left_child->right!= tree->nil) {
        left_child->right->parent = node;
    }
    left_child->parent = node->parent;
    if (node->parent == tree->nil) {
        tree->root = left_child;
    } else if (node == node->parent->right) {
        node->parent->right = left_child;
    } else {
        node->parent->left = left_child;
    }
    left_child->right = node;
    node->parent = left_child;
}

// 插入修复操作
void insert_fixup(rb_tree *tree, rb_node *new_node) {
    while (new_node->parent->color == RED) {
        if (new_node->parent == new_node->parent->parent->left) {
            rb_node *uncle = new_node->parent->parent->right;
            if (uncle->color == RED) {
                new_node->parent->color = BLACK;
                uncle->color = BLACK;
                new_node->parent->parent->color = RED;
                new_node = new_node->parent->parent;
            } else {
                if (new_node == new_node->parent->right) {
                    new_node = new_node->parent;
                    left_rotate(tree, new_node);
                }
                new_node->parent->color = BLACK;
                new_node->parent->parent->color = RED;
                right_rotate(tree, new_node->parent->parent);
            }
        } else {
            rb_node *uncle = new_node->parent->parent->left;
            if (uncle->color == RED) {
                new_node->parent->color = BLACK;
                uncle->color = BLACK;
                new_node->parent->parent->color = RED;
                new_node = new_node->parent->parent;
            } else {
                if (new_node == new_node->parent->left) {
                    new_node = new_node->parent;
                    right_rotate(tree, new_node);
                }
                new_node->parent->color = BLACK;
                new_node->parent->parent->color = RED;
                left_rotate(tree, new_node->parent->parent);
            }
        }
    }
    tree->root->color = BLACK;
}

// 插入节点
void insert(rb_tree *tree, int key) {
    rb_node *new_node = create_rb_node(key);
    rb_node *current = tree->root;
    rb_node *parent = tree->nil;
    while (current!= tree->nil) {
        parent = current;
        if (new_node->key < current->key) {
            current = current->left;
        } else {
            current = current->right;
        }
    }
    new_node->parent = parent;
    if (parent == tree->nil) {
        tree->root = new_node;
    } else if (new_node->key < parent->key) {
        parent->left = new_node;
    } else {
        parent->right = new_node;
    }
    new_node->left = tree->nil;
    new_node->right = tree->nil;
    new_node->color = RED;
    insert_fixup(tree, new_node);
}

// 查找节点
rb_node * search(const rb_tree *tree, int key) {
    rb_node *current = tree->root;
    while (current!= tree->nil && current->key!= key) {
        if (key < current->key) {
            current = current->left;
        } else {
            current = current->right;
        }
    }
    return current;
}

// 查找最小节点
rb_node * minimum(const rb_tree *tree, rb_node *node) {
    while (node->left!= tree->nil) {
        node = node->left;
    }
    return node;
}

// 查找最大节点
rb_node * maximum(const rb_tree *tree, rb_node *node) {
    while (node->right!= tree->nil) {
        node = node->right;
    }
    return node;
}

// 查找后继节点
rb_node * successor(const rb_tree *tree, rb_node *node) {
    if (node->right!= tree->nil) {
        return minimum(tree, node->right);
    }
    rb_node *parent = node->parent;
    while (parent!= tree->nil && node == parent->right) {
        node = parent;
        parent = parent->parent;
    }
    return parent;
}

// 查找前驱节点
rb_node * predecessor(rb_tree *tree, rb_node *node) {
    if (node->left!= tree->nil) {
        return maximum(tree, node->left);
    }
    rb_node *parent = node->parent;
    while (parent!= tree->nil && node == parent->left) {
        node = parent;
        parent = parent->parent;
    }
    return parent;
}
/**
 * 替换红黑树中的一个节点。
 * 
 * @param tree 指向红黑树的指针。
 * @param original 要被替换的节点。
 * @param replacement 用于替换的节点。
 */
void transplant(rb_tree *tree, rb_node *original, rb_node *replacement){
    // 如果要替换的节点是根节点
    if (original->parent == tree->nil) {
        // 将树的根节点指向替换节点
        tree->root = replacement;
    } 
    // 如果要替换的节点是其父节点的左子节点
    else if (original == replacement->parent->left) {
        // 将其父节点的左子节点指向替换节点
        original->parent->left = replacement;
    } 
    // 如果要替换的节点是其父节点的右子节点
    else {
        // 将其父节点的右子节点指向替换节点
        original->parent->right = replacement;
    }
    // 更新替换节点的父节点指针，指向原节点的父节点
    replacement->parent = original->parent;
}

// 删除修复操作
void delete_fixup(rb_tree *tree, rb_node *node) {
    while (node!= tree->root && node->color == BLACK) {
        if (node == node->parent->left) {
            rb_node *sibling = node->parent->right;
            if (sibling->color == RED) {
                sibling->color = BLACK;
                node->parent->color = RED;
                left_rotate(tree, node->parent);
                sibling = node->parent->right;
            }
            if (sibling->left->color == BLACK && sibling->right->color == BLACK) {
                sibling->color = RED;
                node = node->parent;
            } else {
                if (sibling->right->color == BLACK) {
                    sibling->left->color = BLACK;
                    sibling->color = RED;
                    right_rotate(tree, sibling);
                    sibling = node->parent->right;
                }
                sibling->color = node->parent->color;
                node->parent->color = BLACK;
                sibling->right->color = BLACK;
                left_rotate(tree, node->parent);
                node = tree->root;
            }
        } else {
            rb_node *sibling = node->parent->left;
            if (sibling->color == RED) {
                sibling->color = BLACK;
                node->parent->color = RED;
                right_rotate(tree, node->parent);
                sibling = node->parent->left;
            }
            if (sibling->right->color == BLACK && sibling->left->color == BLACK) {
                sibling->color = RED;
                node = node->parent;
            } else {
                if (sibling->left->color == BLACK) {
                    sibling->right->color = BLACK;
                    sibling->color = RED;
                    left_rotate(tree, sibling);
                    sibling = node->parent->left;
                }
                sibling->color = node->parent->color;
                node->parent->color = BLACK;
                sibling->left->color = BLACK;
                right_rotate(tree, node->parent);
                node = tree->root;
            }
        }
    }
    node->color = BLACK;
}

// 删除节点
void delete_node(rb_tree *tree, rb_node *node_to_delete) {
    Color original_color = node_to_delete->color;
    rb_node *child;
    if (node_to_delete->left == tree->nil) {
        child = node_to_delete->right;
        transplant(tree, node_to_delete, child);
    } else if (node_to_delete->right == tree->nil) {
        child = node_to_delete->left;
        transplant(tree, node_to_delete, child);
    } else {
        rb_node *successor_node = minimum(tree, node_to_delete->right);
        original_color = successor_node->color;
        child = successor_node->right;
        if (successor_node->parent == node_to_delete) {
            child->parent = successor_node;
        } else {
            transplant(tree, successor_node, child);
            successor_node->right = node_to_delete->right;
            successor_node->right->parent = successor_node;
        }
        transplant(tree, node_to_delete, successor_node);
        successor_node->left = node_to_delete->left;
        successor_node->left->parent = successor_node;
        successor_node->color = node_to_delete->color;
    }
    if (original_color == BLACK) {
        delete_fixup(tree, child);
    }
    free(node_to_delete);
}

// 中序遍历红黑树
void inorder_traversal(rb_node *node) {
    if (node!= NULL) {
        inorder_traversal(node->left);
        printf("%d ", node->key);
        inorder_traversal(node->right);
    }
}

// 销毁红黑树
void destroy_rb_tree(rb_tree *tree) {
    if (tree == NULL) {
        return;
    }
    // 递归释放所有节点
    destroy_rb_tree_helper(tree->root);
    // 释放哨兵节点
    free(tree->nil);
    // 释放红黑树结构体
    free(tree);
}

// 递归释放节点
void destroy_rb_tree_helper(rb_node *node) {
    if (node == NULL) {
        return;
    }
    // 递归释放左子树
    destroy_rb_tree_helper(node->left);
    // 递归释放右子树
    destroy_rb_tree_helper(node->right);
    // 释放当前节点
    free(node);
}